# coding=utf-8

# python pre_work_to_colmap_format.py 1403715300
# 1403715302712143104.png ~ 1403715307762142976.png | 590~691
# 1403715394662142976.png ~ 1403715400362142976.png | 2429~2543

from lib2to3.pgen2.token import NAME
import os
from re import I
import sys
import numpy as np

# 全局变量定义
g_image_width = 0
g_image_height = 0
g_fxycxy = [0,0,0,0]
g_key_frame_cnt = 0

def qvec2rotmat(qvec): # 四元数 q(w,x,y,z) 对应的旋转矩阵
    return np.array([
        [1 - 2 * qvec[2]**2 - 2 * qvec[3]**2,
         2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3],
         2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2]],
        [2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3],
         1 - 2 * qvec[1]**2 - 2 * qvec[3]**2,
         2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1]],
        [2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2],
         2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1],
         1 - 2 * qvec[1]**2 - 2 * qvec[2]**2]])

def findAllFile(path_str):
    onlynames = []
    for root, ds, fs in os.walk(path_str):
        for f in fs:
            onlynames.append(f)
    return onlynames

def point3d_to_point2d(point3d, qvec, tvec):
    # todo : 单位还没确认
    # 基于旋转和平移的外参矩阵
    e = np.zeros((4, 4))
    e[:3, :3] = qvec2rotmat(qvec)
    e[:3, 3] = tvec
    e[3, 3] = 1
    extrinsic = e
    # 内参矩阵 g_fxycxy
    fx, fy, cx, cy = g_fxycxy[0], g_fxycxy[1], g_fxycxy[2], g_fxycxy[3]
    i = np.array([
        [fx, 0, cx],
        [0, fy, cy],
        [0, 0, 1]
    ], dtype = np.float32)
    intrinsic = i
    # 基于内参矩阵和外参矩阵将世界坐标系中的3D点转到图像坐标系中
    P3D = np.array([
        [point3d[0]],
        [point3d[1]],
        [point3d[2]],
        [1]
    ], dtype = np.float32)
    point2d_cam = np.dot(extrinsic,P3D)
    point2d = np.dot(intrinsic,point2d_cam[0:3]) # (3,3)
    point2d = point2d / point2d[2]
    u_img = point2d[0][0]
    v_img = point2d[1][0]
    return [u_img, v_img]

def worker_func_read_cameras_config():
    root_path = os.getcwd()
    sfm_reader = open(root_path + '/sfm.txt', 'r')
    for i in range(3):
        # print(sfm_reader.readline().strip())
        line_elements = sfm_reader.readline().strip().split()
        # 第一行 图像大小
        if i == 0:
            g_image_width_str = line_elements[1]
            g_image_height_str = line_elements[2]
            global g_image_width 
            g_image_width = int(g_image_width_str)
            global g_image_height
            g_image_height = int(g_image_height_str)
        # 第二行 关键帧数目
        if i == 1:
            global g_key_frame_cnt 
            g_key_frame_cnt = int(line_elements[0])
        # 第三行 获取内参 fx fy cx cy
        if i == 2:
            fx_str = line_elements[1]
            fy_str = line_elements[2]
            cx_str = line_elements[3]
            cy_str = line_elements[4]
            global g_fxycxy 
            g_fxycxy = map(float, line_elements[1:5])
            fxycxy_str = fx_str + ' ' + fy_str + ' ' + cx_str + ' ' + cy_str
            # 打开并准备写入 cameras.txt 
            with open(root_path + '/cameras.txt', 'w') as f_colmap_cameras_writer:
                f_colmap_cameras_writer.write("# Camera list with one line of data per camera:\n")
                f_colmap_cameras_writer.write('# CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[fx,fy,cx,cy]\n')
                f_colmap_cameras_writer.write('# Number of cameras: 1\n')
                camera_id_str = '1'
                camera_model_str = 'PINHOLE'
                line_to_write = camera_id_str + ' ' + camera_model_str + ' ' + g_image_width_str + ' ' + g_image_height_str + ' ' + fxycxy_str + '\n'
                f_colmap_cameras_writer.write(line_to_write)
            break


def worker_func_read_images_pose_point2d(init_time_1e6):
    root_path = os.getcwd()
    sfm_reader = open(root_path + '/sfm.txt', 'r')
    lines = sfm_reader.readlines()
    # 数据缓存变量准备
    IMAGE_ID = []
    QW, QX, QY, QZ = [], [], [], []
    TX, TY, TZ = [], [], []
    TIMESTAMP, NAME = [], []
    POINT3D = []
    POINT3D_IN_FRAMES = []
    # 开启循环遍历，将数据保存至缓存变量中
    for i in range(len(lines)):
        # 第一行跳过
        if i == 0: 
            continue
        # 第二行是序号
        if i == 1: 
            # print(lines[i])
            continue
        # 第三行开始 到 所有位姿全部遍历
        if i >= 2 and i < (g_key_frame_cnt+2): 
            # 将所有的位姿存储起来 
            one_line_elements_str = lines[i].split()
            IMAGE_ID.append(int(one_line_elements_str[0]))
            QW.append(float(one_line_elements_str[5]))
            QX.append(float(one_line_elements_str[6]))
            QY.append(float(one_line_elements_str[7]))
            QZ.append(float(one_line_elements_str[8]))
            TX.append(float(one_line_elements_str[9]))
            TY.append(float(one_line_elements_str[10]))
            TZ.append(float(one_line_elements_str[11]))
            TIMESTAMP.append(float(one_line_elements_str[12]) - init_time_1e6)
            # 根据 timestamp 确定照片的名字
            cam0_data_path = '/home/mxg/MySpace/catkin_ws/DataSets/Euroc/V1_01_easy/mav0/cam0/data/'
            image_names = findAllFile(cam0_data_path)
            for image_name in image_names:
                image_name_num = int(int(int(image_name.strip('.png'))*1e-3) - init_time_1e6*1e6)
                timestamp = int(TIMESTAMP[-1]*1E6)
                if abs(image_name_num - timestamp)<=1: # 中间可能会有 四舍五入 的情况
                    # print(image_name)
                    NAME.append(image_name)
                    break
        # 读取 特征点的个数 
        if i == (g_key_frame_cnt+2):
            print('Total 3D points: ' + lines[i])
        # 开始遍历所有三维特征点
        if i > (g_key_frame_cnt+2):
            one_line_elements_str = lines[i].split()
            POINT3D.append(map(float,one_line_elements_str[0:3])) # 在这里隔了一个数字，表示帧的数目 0～2 4～end
            POINT3D_IN_FRAMES.append(map(int,one_line_elements_str[4:]))
    
    # print(len(IMAGE_ID))
    # print(IMAGE_ID[0])
    # 构建 images.txt 和 points3D.txt 文件
    with open(root_path + '/images.txt', 'w') as f_colmap_images_writer, open(root_path + '/points3D.txt', 'w') as f_colmap_points3D_writer:
        
        # images.txt 文件头
        f_colmap_images_writer.write('# Image list with two lines of data per image:\n')
        f_colmap_images_writer.write('#   IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n')
        f_colmap_images_writer.write('#   POINTS2D[] as (X, Y, POINT3D_ID)\n')
        f_colmap_images_writer.write('# Number of images: TXT, mean observations per image: TXT\n')
        
        # points3D.txt 文件头
        f_colmap_points3D_writer.write('# 3D point list with one line of data per point:\n')
        f_colmap_points3D_writer.write('#   POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)\n')
        f_colmap_points3D_writer.write('# Number of points: 6145, mean track length: 10.593816110659072\n')
        
        # 开始填充 images.txt 内容
        print('# 开始填充 images.txt 内容')
        for i in range(len(IMAGE_ID)): 
            image_id = str(IMAGE_ID[i]) 
            qw, qx, qy, qz = str(QW[i]), str(QX[i]), str(QY[i]), str(QZ[i])
            tx, ty, tz = str(TX[i]), str(TY[i]), str(TZ[i])
            camera_id = str(1) # 只有一个相机
            name = str(NAME[i])
            # 写入每一帧的位姿信息
            f_colmap_images_writer.write(image_id + ' ' + qw + ' ' + qx + ' '  + qy + ' ' + qz + ' ' + tx + ' ' + ty + ' ' + tz + ' ' + camera_id + ' ' + name + '\n')
            # 写入每一帧图像看到的2D特征点 & 特征点在三维空间中的ID
            POINTS2D = ""
            for j in range(len(POINT3D)): # j 是 3D 点的索引
                # 遍历当前  POINT3D_IN_FRAMES ，如果此帧的 ID 存在于 POINT3D 所属，则将其投影到图像的上，得到 point2d
                if i in POINT3D_IN_FRAMES[j]:
                    # 3D 特征点投影至 2D
                    point2d = point3d_to_point2d(POINT3D[j], [float(qw), float(qx), float(qy), float(qz)], [float(tx), float(ty), float(tz)])
                    POINTS2D += str(point2d[0]) + ' ' + str(point2d[1]) + ' ' + str(j) + ' '
            f_colmap_images_writer.write(POINTS2D + '\n')
    
        # 开始填充 points3D.txt 内容
        print('# 开始填充 points3D.txt 内容')
        TRACK = []
        for i in range(len(POINT3D)): 
            point3d_id = str(i) # 从0开始
            point3d_x, point3d_y, point3d_z = str(POINT3D[i][0]), str(POINT3D[i][1]), str(POINT3D[i][2])
            point3d_r, point3d_g, point3d_b = str(122), str(122), str(122)
            point3d_error = str(0.25)
            f_colmap_points3D_writer.write(point3d_id + ' ' + point3d_x + ' ' + point3d_y + ' ' + point3d_z + ' ' + point3d_r + ' ' + point3d_g + ' ' + point3d_b + ' ' + point3d_error + ' ')
            # 将上述内容写入到 points3D.txt 中
            # todo: TRACK[] as (IMAGE_ID, POINT2D_IDX), TMD !!!
            # 遍历 image ，查看 第i个POINT3D 是否在 第j个image 上，并统计 这个 POINT3D 是image上的第几个 
            for j in range(len(IMAGE_ID)):
                # 判断 第i个POINT3D 是否在 第j个image 上
                if j in POINT3D_IN_FRAMES[i]:
                    # 记录 image_id
                    image_id = j
                    # 统计 第i个POINT3D 是 第j个image 上第几个点
                    # 先设定一个最小的序号 0, 其中 [x,x] 中的 image_id 肯定是对的，但是不确定序号是多少
                    track = [image_id, 0] 
                    while track in TRACK: # 如果上面的 track 已经存在过，则将序号加 1 ，直到不存在
                        track[1] += 1 
                    # 将 track 放入 TRACK 中
                    TRACK.append(track)
                    # 而后将 track 输出到 points3D.txt 中
                    f_colmap_points3D_writer.write(str(track[0]) + ' ' + str(track[1]) + ' ')
            f_colmap_points3D_writer.write('\n')
            if i % 10 == 0:
                print('Doing %d point3d, total : %d' % (i, len(POINT3D)))


def main(init_time_1e6):
    print('Current Path : ' + os.getcwd())
    # 判断是否存在 sfm.txt 文件
    if os.path.exists('./sfm.txt') is False:
        raise NameError('Can not find sfm.txt !!!')
    
    # 读取相机内参，并将其写入 cameras.txt 中
    worker_func_read_cameras_config()
    print("Camera Intrinsics: ", g_fxycxy)
    print("(The size of image: %d x %d)" % (g_image_width, g_image_height))
    print("(Key frame count: %d)" % g_key_frame_cnt)

    # 读取每一帧关键帧的位姿信息，并将每一帧的位姿 CAMERA_ID IMG_NAME POINTS2D[]
    # 写入到 images.txt 中
    worker_func_read_images_pose_point2d(init_time_1e6)

if __name__ == '__main__':
    init_time_1e6 = int(sys.argv[1])
    main(init_time_1e6)
